#include "pso.h"
#include "velocity.h"
#include "position.h"
#include "rand.h"
#include "params.h"
#include <iomanip>


//-------------------------------------------------------------------------
// Public

Pso::Pso(Params *p_params, bool p_log_stats/* = true*/)
    : m_params(p_params)
    , m_log_stats(p_log_stats)
{
    m_best_overall_solution = new Particle(this);
}

Pso::~Pso()
{
    // Delete the best overall solution.
    delete m_best_overall_solution;
}

double Pso::exec()
{
    double solution = run();

    return solution;
}

double Pso::exec(double p_alpha_1, double p_alpha_2)
{
    double solution = 0;

    m_params->alpha_1 = p_alpha_1;
    m_params->alpha_2 = p_alpha_2;

    solution = run();

    return solution;
}

double Pso::execOmega(double p_omega_1, double p_omega_2)
{
    double solution = 0;

    m_params->omega_top = p_omega_1;
    m_params->omega_bottom = p_omega_2;

    solution = run();

    return solution;
}

const Params *Pso::getParams()
{
    return m_params;
}

//-------------------------------------------------------------------------
// Private

double Pso::run()
{
    int particle_count;
    int iterations = m_params->iterations;
    double omega_top = m_params->omega_top;
    double omega_bottom = m_params->omega_bottom;

    // Open the logs where we will store stats.
    if(m_log_stats) {
        openLogs();
    }
    
    // Create the initial particles.
    initSwarm();
    
    // Init the overall best for the first time.
    updateOverallBest(true);

    particle_count = m_particles.size();

    // Start the algorithm.
    for(int i = 0; i < iterations; ++i) {
        // Update the omega value for this iteration.
        m_current_omega = omega_top - ((static_cast<double>(i) / iterations) * 
            (omega_top - omega_bottom));

        // Update particles' velocity and position.
        for(int j = 0; j < particle_count; ++j) {
            updateVelocity(j);
            updatePosition(j);
        }

        // Update the best solution found so far.
        updateOverallBest();

        // And gather statistics for the plots.
        if(m_log_stats) {
            collectStats(i);
        }
    }

    // We no longer need to write into log files.
    if(m_log_stats) {
        closeLogs();
    }

    // The swarm is no longer required.
    destroySwarm();

    return m_best_overall_solution->getFitness();
}

void Pso::initSwarm()
{
    Particle *particle;

    // Create the required ammount of particles. They are created
    // ... randomly by nature.
    for(unsigned int i = 0; i < m_params->swarm_size; ++i) {
        particle = new Particle(this);
        // And calculate the initial fitness given the random values.
        particle->calculateFitness();
        m_particles.push_back(particle);
    }
}

void Pso::destroySwarm() 
{
    unsigned int swarm_size = m_particles.size();

    // Delete the particles.
    for (unsigned int i = 0; i < swarm_size; ++i) {
        delete m_particles.at(i);
    }
    m_particles.clear();
}

void Pso::updateVelocity(unsigned int p_particle_i)
{
    Particle *particle = m_particles.at(p_particle_i);
    int dimensions = m_params->dimensions;
    double inertia;
    double cognitive;
    double social;
    double historical;
    const Velocity *current_velocity = particle->getVelocity();
    const Position *current_position = particle->getPosition();
    const Position *best_position = particle->getBestPosition();
    const Position *global_best_position = getBestNeighbourPosition(p_particle_i);
    const Position *best_solution_found = m_best_overall_solution->getPosition();
    Velocity new_velocity(m_params->dimensions, 0.0);

    // Update each component individually.
    for(int i = 0; i < dimensions; ++i) {
        // Inertia component.
        inertia = m_current_omega * current_velocity->getComponent(i);

        // Cognitive component.
        cognitive = (m_params->alpha_1 * randy::random_uniform()) *
            (best_position->getComponent(i) - current_position->getComponent(i));

        // Social component.
        social = (m_params->alpha_2 * randy::random_uniform()) *
            (global_best_position->getComponent(i) - current_position->getComponent(i));

        // Historical component.
        historical = (m_params->alpha_3 * randy::random_uniform()) *
            (best_solution_found->getComponent(i) - current_position->getComponent(i));

        // Calculate resulting component.
        new_velocity.setComponent(i, inertia + cognitive + social + historical);

        // We are not bounding the velocity component to a maximum value.
    }

    particle->setVelocity(new_velocity);
}

void Pso::updatePosition(unsigned int p_particle_i)
{
    Particle *particle = m_particles.at(p_particle_i);
    int dimensions = m_params->dimensions;
    const Velocity *current_velocity = particle->getVelocity();
    const Position *current_position = particle->getPosition();    
    Position new_position(m_params->dimensions, 0.0);
    
    // Update each component individually.
    for(int i = 0; i < dimensions; ++i) {
        new_position.setComponent(i, current_position->getComponent(i) +
            current_velocity->getComponent(i));

        // Not bounding the resulting component to a maximum value.
    }

    particle->setPosition(new_position);
}

void Pso::updateOverallBest(bool first_update/* = false*/)
{
    int particle_count = m_particles.size();
    int i = 0;

    // Get the initial value to test for improvement in a different
    // ... way if this is the first update for the particles.

    if(first_update == true) {
        // Arbitrarily define the first particle as the best.
        (*m_best_overall_solution) = *m_particles.at(0);
        // Start from particle two.
        i = 1;
    }

    // Find if a particle has improved abover our recorded best.
    for(; i < particle_count; ++i) {
        // Minimization going on here.
        if(m_particles.at(i)->getFitness() < m_best_overall_solution->getFitness()) {
            // This particle rocks!
            (*m_best_overall_solution) = *m_particles.at(i);
        }
    }
}

const Position *Pso::getBestNeighbourPosition(unsigned int p_particle_i)
{
    int neighbour_count = m_params->neighbours;
    unsigned int next_particle_i;
    Particle *current_particle = m_particles.at(p_particle_i);
    const Position *global_best_position = current_particle->getPosition();
    double global_best_fitness = current_particle->getFitness();

    // Examine the neighbours of this particle for the best position
    // ... (the one with best fitness).
    for(int i = 1; i < neighbour_count; ++i) {
        next_particle_i = (p_particle_i + i) % neighbour_count;
        current_particle = m_particles.at(next_particle_i);

        // We are minimizing.
        if(current_particle->getFitness() < global_best_fitness) {
            // Update the global best with this particle.
            global_best_fitness = current_particle->getFitness();
            global_best_position = current_particle->getPosition();
        }
    }

    return global_best_position;
}

void Pso::openLogs()
{
    m_logger.open("particle_1.plot");
    m_logger.open("particle_2.plot");
    //m_logger.open("particle_3.plot");
    m_logger.open("solution.plot");
}

void Pso::closeLogs()
{
    m_logger.close();
}

void Pso::collectStats(unsigned int p_iteration)
{
    // Log all the dimensions of particle 1.
    m_logger.log("particle_1.plot").precision(8);
    m_logger.log("particle_1.plot") << std::setw(3) << p_iteration;
    for(unsigned int i = 0; i < m_params->dimensions; ++i) {
        m_logger.log("particle_1.plot") 
            << std::setw(20)
            << m_particles.at(0)->getPosition()->getComponent(i);
    }
    m_logger.log("particle_1.plot") << std::endl;


    // Log all the dimensions of particle 2.
    m_logger.log("particle_2.plot").precision(8);
    m_logger.log("particle_2.plot")  << std::setw(3) << p_iteration;
    for(unsigned int i = 0; i < m_params->dimensions; ++i) {
        m_logger.log("particle_2.plot") 
            << std::setw(20)
            << m_particles.at(9)->getPosition()->getComponent(i);
    }
    m_logger.log("particle_2.plot") << std::endl;


    // Log all the dimensions of the solution plus fitness.
    m_logger.log("solution.plot").precision(8);
    m_logger.log("solution.plot") << std::setw(3) << p_iteration;
    for(unsigned int i = 0; i < m_params->dimensions; ++i) {
        m_logger.log("solution.plot") 
            << std::setw(20)
            << m_best_overall_solution->getPosition()->getComponent(i);
    }
    m_logger.log("solution.plot") 
        << std::setw(20) 
        << m_best_overall_solution->getFitness() 
        << std::endl;
}
